//
// Created by whitby on 8/22/23.
//

#include "Optimize.h"
#include "G2oTypes.h"
#include "Sensor/Camera.h"
#include "Utils/Converter.h"

#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/robust_kernel_impl.h>

using namespace std;

namespace mono_orb_slam2 {

    void Optimize::initialOptimize(const std::shared_ptr<KeyFrame> &lastKeyFrame,
                                   const std::shared_ptr<KeyFrame> &curKeyFrame) {
        int iteration = 20;

        // 1. set-up optimizer
        g2o::SparseOptimizer optimizer;
        auto *solver = new g2o::OptimizationAlgorithmLevenberg(
                g2o::make_unique<g2o::BlockSolver_6_3>(g2o::make_unique<LinearSolverEigen_6_3>())
        );
        optimizer.setAlgorithm(solver);

        // 2. set vertices and edges
        Pose T1w = lastKeyFrame->getPose();
        auto *vPose1 = new VertexPose(T1w.R, T1w.t);
        vPose1->setId(0);
        vPose1->setFixed(true);
        optimizer.addVertex(vPose1);

        Pose T2w = curKeyFrame->getPose();
        auto *vPose2 = new VertexPose(T2w.R, T2w.t);
        vPose2->setId(1);
        optimizer.addVertex(vPose2);

        int mpId = 2;
        vector<shared_ptr<MapPoint>> mapPoints = curKeyFrame->getTrackedMapPoints();
        for (const auto &mp: mapPoints) {
            auto *vPoint = new VertexPoint(mp->getPos());
            vPoint->setId(mpId++);
            vPoint->setMarginalized(true);
            optimizer.addVertex(vPoint);

            unsigned int idx1 = mp->getFeatureId(lastKeyFrame);
            const cv::KeyPoint &kp1 = lastKeyFrame->key_points[idx1];
            auto *eVisual1 = new EdgeMono();
            eVisual1->setVertex(0, vPoint);
            eVisual1->setVertex(1, vPose1);
            eVisual1->setMeasurement({kp1.pt.x, kp1.pt.y});
            eVisual1->setInformation(Eigen::Matrix2d::Identity());

            auto *rk1 = new g2o::RobustKernelHuber;
            eVisual1->setRobustKernel(rk1);
            rk1->setDelta(5.991);

            optimizer.addEdge(eVisual1);

            unsigned int idx2 = mp->getFeatureId(curKeyFrame);
            const cv::KeyPoint &kp2 = curKeyFrame->key_points[idx2];
            auto *eVisual2 = new EdgeMono();
            eVisual2->setVertex(0, vPoint);
            eVisual2->setVertex(1, vPose2);
            eVisual2->setMeasurement({kp2.pt.x, kp2.pt.y});
            eVisual2->setInformation(Eigen::Matrix2d::Identity());

            auto *rk2 = new g2o::RobustKernelHuber;
            eVisual2->setRobustKernel(rk2);
            rk2->setDelta(5.991);

            optimizer.addEdge(eVisual2);
        }

        // 3. optimize
        optimizer.initializeOptimization();
        optimizer.optimize(iteration);

        // 4. recover
        auto pose2 = vPose2->estimate();
        curKeyFrame->setPose(pose2.rotation().toRotationMatrix().cast<float>(), pose2.translation().cast<float>());

        mpId = 2;
        for (auto &mp: mapPoints) {
            auto *vPoint = dynamic_cast<VertexPoint *>(optimizer.vertex(mpId++));
            mp->setPos(vPoint->estimate().cast<float>());
            mp->update();
        }
    }

    int Optimize::poseOptimize(const std::shared_ptr<Frame> &frame) {
        // 1. set optimizer
        g2o::SparseOptimizer optimizer;
        auto *solver = new g2o::OptimizationAlgorithmLevenberg(
                g2o::make_unique<g2o::BlockSolver_6_3>(g2o::make_unique<LinearSolverEigen_6_3>())
        );
        optimizer.setAlgorithm(solver);

        int numInitialCorrespondences = 0;

        // 2. set vertex and edges
        const Pose Tcw = frame->T_cw;
        auto *vPose = new VertexPose(Tcw.R, Tcw.t);
        vPose->setId(0);
        optimizer.addVertex(vPose);

        // set map-point vertices
        const int n = frame->num_kps;
        vector<EdgeMonoPose *> vecEdge;
        vector<size_t> vecIndices;
        vecEdge.reserve(n), vecIndices.reserve(n);

        const float deltaMono = sqrtf(5.991);

        {
            lock_guard<mutex> lock(MapPoint::global_mutex);
            for (int i = 0; i < n; ++i) {
                const shared_ptr<MapPoint> &mp = frame->map_points[i];
                if (mp != nullptr && !mp->isBad()) {
                    numInitialCorrespondences++;
                    const cv::KeyPoint &kp = frame->key_points[i];

                    auto *eVisual = new EdgeMonoPose(mp->getPos());
                    eVisual->setVertex(0, vPose);
                    eVisual->setMeasurement({kp.pt.x, kp.pt.y});
                    const float invSigma2 = ORBExtractor::getInvSquareSigma(kp.octave);
                    eVisual->setInformation(invSigma2 * Eigen::Matrix2d::Identity());

                    auto *rk = new g2o::RobustKernelHuber;
                    eVisual->setRobustKernel(rk);
                    rk->setDelta(deltaMono);

                    optimizer.addEdge(eVisual);
                    vecEdge.push_back(eVisual);
                    vecIndices.push_back(i);
                }
            }
        }

        if (numInitialCorrespondences < 3) return 0;

        // 3. optimize
        // we perform 4 optimizations, after each optimization we classify observation as inlier/outlier
        // at the next optimization, outliers are not included, but they can be optimized again at the end.
        const int iters[4] = {10, 10, 10, 10};

        vector<bool> vecBeInlier(numInitialCorrespondences, true);
        for (auto iter: iters) {
            vPose->setEstimate({Tcw.R.cast<double>(), Tcw.t.cast<double>()});

            optimizer.initializeOptimization(0);
            optimizer.optimize(iter);

            for (unsigned i = 0; i < numInitialCorrespondences; ++i) {
                auto *e = vecEdge[i];

                if (!vecBeInlier[i]) e->computeError();

                const double chi2 = e->chi2();
                if (chi2 > 5.991) {
                    vecBeInlier[i] = false;
                    e->setLevel(1);
                } else {
                    vecBeInlier[i] = true;
                    e->setLevel(0);
                }

                if (iter == 2)
                    e->setRobustKernel(nullptr);
            }
        }

        // 4. recover
        auto optimizedPose = vPose->estimate();
        frame->setPose({optimizedPose.rotation().toRotationMatrix().cast<float>(), optimizedPose.translation().cast<float>()});

        // discard outlier map-points
        int numInlier = 0;
        for (int i = 0; i < numInitialCorrespondences; ++i) {
            if (!vecBeInlier[i]) {
                shared_ptr<MapPoint> &mp = frame->map_points[vecIndices[i]];
                mp->track_in_view = false;
                mp->last_frame_seen = frame->id;
                mp = nullptr;
            } else {
                numInlier++;
            }
        }

        return numInlier;
    }

    void Optimize::localBundleAdjustment(const std::shared_ptr<KeyFrame> &curKeyFrame, Map *pointMap, bool *stopFlag) {
        // 1. extract local map
        // local keyframes: first breath search from current keyframe
        list<shared_ptr<KeyFrame>> localKeyFrames;
        localKeyFrames.push_back(curKeyFrame);
        unsigned int curKFId = curKeyFrame->id;
        curKeyFrame->BA_local_for_kf = curKFId;

        const vector<shared_ptr<KeyFrame>> neighKeyFrames = curKeyFrame->getConnectedKFs();
        for (const auto &kf: neighKeyFrames) {
            kf->BA_local_for_kf = curKFId;
            localKeyFrames.push_back(kf);
        }

        // local map-points seen in local keyframes
        list<shared_ptr<MapPoint>> localMapPoints;
        for (const auto &kf: localKeyFrames) {
            vector<shared_ptr<MapPoint>> mapPoints = kf->getMapPoints();
            for (const auto &mp: mapPoints) {
                if (mp && mp->BA_local_for_kf != curKFId) {
                    localMapPoints.push_back(mp);
                    mp->BA_local_for_kf = curKFId;
                }
            }
        }

        // fixed keyframe (see local map-points but not in local-keyframes)
        list<shared_ptr<KeyFrame>> fixedKeyFrames;
        for (const auto &mp: localMapPoints) {
            map<shared_ptr<KeyFrame>, size_t> observations = mp->getObservations();
            for (const auto &obsPair: observations) {
                const shared_ptr<KeyFrame> &kf = obsPair.first;
                if (kf->BA_local_for_kf != curKFId && kf->BA_fixed_for_kf != curKFId) {
                    kf->BA_fixed_for_kf = curKFId;
                    fixedKeyFrames.push_back(kf);
                }
            }
        }

        // 2. set optimizer
        g2o::SparseOptimizer optimizer;
        auto *solver = new g2o::OptimizationAlgorithmLevenberg(
                g2o::make_unique<g2o::BlockSolver_6_3>(g2o::make_unique<LinearSolverEigen_6_3>())
        );

        optimizer.setAlgorithm(solver);

        if (stopFlag != nullptr)
            optimizer.setForceStopFlag(stopFlag);

        // 3. set vertices and edges
        // local keyframe vertices
        unsigned long maxKFId = 0;
        for (const auto &kf: localKeyFrames) {
            const Pose Tcw = kf->getPose();
            auto *vPose = new VertexPose(Tcw.R, Tcw.t);
            vPose->setId(kf->id);
            vPose->setFixed(kf->id == 0);
            optimizer.addVertex(vPose);

            maxKFId = max(maxKFId, kf->id);
        }

        // fixed keyframe vertices
        for (const auto &kf: fixedKeyFrames) {
            const Pose Tcw = kf->getPose();
            auto *vPose = new VertexPose(Tcw.R, Tcw.t);
            vPose->setId(kf->id);
            vPose->setFixed(true);
            optimizer.addVertex(vPose);

            maxKFId = max(maxKFId, kf->id);
        }

        unsigned int expectSize = localKeyFrames.size() * localMapPoints.size();
        vector<VertexPoint *> vecPoints;
        vecPoints.reserve(localMapPoints.size());

        vector<EdgeMono *> edges;
        edges.reserve(expectSize);

        vector<pair<shared_ptr<KeyFrame>, shared_ptr<MapPoint>>> vecEdgeObs;
        vecEdgeObs.reserve(expectSize);

        const float thHuberMono = sqrtf(5.991);
        int numEdge = 0;
        int mpId = maxKFId + 1;
        for (const auto &mp: localMapPoints) {
            auto *vPoint = new VertexPoint(mp->getPos());
            vPoint->setId(mpId++);
            vPoint->setMarginalized(true);
            optimizer.addVertex(vPoint);
            vecPoints.push_back(vPoint);

            const map<shared_ptr<KeyFrame>, size_t> observation = mp->getObservations();
            for (const auto &obsPair: observation) {
                const shared_ptr<KeyFrame> &kf = obsPair.first;
                const cv::KeyPoint &kp = kf->key_points[obsPair.second];

                auto *eVisual = new EdgeMono();
                eVisual->setVertex(0, vPoint);
                eVisual->setVertex(1, optimizer.vertex(kf->id));
                eVisual->setMeasurement({kp.pt.x, kp.pt.y});
                const float invSigma2 = ORBExtractor::getInvSquareSigma(kp.octave);
                eVisual->setInformation(invSigma2 * Eigen::Matrix2d::Identity());

                auto *rk = new g2o::RobustKernelHuber;
                eVisual->setRobustKernel(rk);
                rk->setDelta(thHuberMono);

                optimizer.addEdge(eVisual);
                edges.push_back(eVisual);
                vecEdgeObs.emplace_back(kf, mp);
                numEdge++;
            }
        }

        // 4. optimize
        optimizer.initializeOptimization();
        optimizer.optimize(5);

        bool doMore = true;
        if (stopFlag != nullptr && *stopFlag) doMore = false;

        if (doMore) {
            // 5. check inlier observations
            for (int i = 0; i < numEdge; ++i) {
                auto *e = edges[i];
                if (e->chi2() > 5.991 || !e->isDepthPositive()) {
                    e->setLevel(1);
                }

                e->setRobustKernel(nullptr);
            }

            // 6. optimize again without outliers
            optimizer.initializeOptimization(0);
            optimizer.optimize(10);
        }

        // 7. check inlier observations, discard outliers
        vector<pair<shared_ptr<KeyFrame>, shared_ptr<MapPoint>>> vecToErase;
        vecToErase.reserve(numEdge / 2);
        for (int i = 0; i < numEdge; ++i) {
            auto *e = edges[i];
            if (e->chi2() > 5.991 || !e->isDepthPositive()) {
                vecToErase.push_back(vecEdgeObs[i]);
            }
        }

        // get map update mutex
        lock_guard<mutex> lock(pointMap->map_update_mutex);

        for (const auto &erasePair: vecToErase) {
            const shared_ptr<KeyFrame> &kf = erasePair.first;
            const shared_ptr<MapPoint> &mp = erasePair.second;

            if (mp->isBad()) continue;
            kf->eraseMapPoint(mp->getFeatureId(kf));
            mp->eraseObservation(kf);
        }

        // 8. recover
        for (const auto &kf: localKeyFrames) {
            auto pose = dynamic_cast<VertexPose *>(optimizer.vertex(kf->id))->estimate();
            const Eigen::Matrix4f optimizeTcw = pose.to_homogeneous_matrix().cast<float>();
            kf->setPose(optimizeTcw.block<3, 3>(0, 0), optimizeTcw.block<3, 1>(0, 3));
        }

        mpId = maxKFId + 1;
        for (const auto &mp: localMapPoints) {
            auto *vPoint = dynamic_cast<VertexPoint *>(optimizer.vertex(mpId++));
            if (!mp->isBad()) {
                mp->setPos(vPoint->estimate().cast<float>());
                mp->update();
            }
        }
    }

} // mono_orb_slam2